"""
Author: Morphlng
Date: 2024-03-22 10:50:28
LastEditTime: 2024-04-16 16:19:48
LastEditors: Morphlng
Description: Adapter for metadrive environment
FilePath: /DrivingGym/src/driving_gym/simulation/adapter/metadrive/metadrive_adapter.py
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from metadrive.component.sensors.base_camera import BaseSensor
    from metadrive.component.vehicle.vehicle_type import BaseVehicle

from driving_gym.misc.util import norm_angle, override
from driving_gym.simulation.adapter_interface import AdapterInterface
from driving_gym.simulation.common.geom import (
    BoundingBox,
    Location,
    Rotation,
    Transform,
    Vector3D,
)


class MetadriveAdapter(AdapterInterface):
    """Environment adapter for metadrive environment"""

    def __init__(self, config: dict):
        super().__init__(config)
        self.env = None
        self.env_config = None
        self._tick_cnt = 0
        self._last_obs = None
        self._last_reward = None
        self._last_terminated = None
        self._last_truncated = None
        self._last_info = None
        self._action_dict = {}
        self._agent_states = {}

    @property
    def agents(self) -> dict[str, "BaseVehicle"]:
        if self.env is None:
            raise ValueError("Environment is not initialized")
        return self.env.agents

    @property
    def sensors(self) -> dict[str, "BaseSensor"]:
        if self.env is None:
            raise ValueError("Environment is not initialized")
        return self.env.engine.sensors

    @property
    def static_world(self):
        if self.env is None:
            raise ValueError("Environment is not initialized")
        return self.env.engine.physics_world.static_world

    @property
    def dynamic_world(self):
        if self.env is None:
            raise ValueError("Environment is not initialized")
        return self.env.engine.physics_world.dynamic_world

    @override(AdapterInterface)
    def initialize(self):
        from metadrive.envs.marl_envs import MultiAgentMetaDrive

        if self.env is not None:
            self.cleanup()
        self.env = MultiAgentMetaDrive(config=self.config)
        self.env_config = self.env.config

    @override(AdapterInterface)
    def cleanup(self):
        if self.env is not None:
            self.env.close()
            self.env = None

        self._last_obs = None
        self._last_reward = None
        self._last_terminated = None
        self._last_truncated = None
        self._last_info = None
        self._action_dict.clear()
        self._agent_states.clear()

    @override(AdapterInterface)
    def get_map(self):
        # TODO: Metadrive is planning to support OpenDrive
        return None

    @override(AdapterInterface)
    def get_transform(self, identifier: str):
        # Using agent_states will lose z-axis information
        # So we use the origin node to get the transform
        agent = self.agents[identifier]
        pos = agent.origin.getPos()
        hpr = agent.origin.getHpr()  # heading(yaw), pitch, roll

        return Transform(
            Location(pos.x, pos.y, pos.z),
            Rotation(
                pitch=norm_angle(hpr.y),
                yaw=norm_angle(hpr.x + 90),
                roll=norm_angle(hpr.z),
            ),
        )

    @override(AdapterInterface)
    def get_velocity(self, identifier: str):
        # Using agent_states will lose z-axis information
        # `agent.body.get_linear_velocity()` returns a 3D vector
        agent = self.agents[identifier]
        velocity = agent.body.get_linear_velocity()
        return Vector3D(velocity.x, velocity.y, velocity.z)

    @override(AdapterInterface)
    def get_acceleration(self, identifier: str):
        # TODO: Metadrive does not provide acceleration vector
        state = self.get_state(identifier)
        return Vector3D(state["throttle_brake"], 0, 0)

    @override(AdapterInterface)
    def get_bounding_box(self, identifier: str):
        state = self.get_state(identifier)
        size = state["size"]
        extent = Vector3D(size[0] / 2, size[1] / 2, size[2] / 2)
        return BoundingBox(Location(), extent)

    @override(AdapterInterface)
    def get_sensor_data(self, sensor_identifier: str):
        from metadrive.component.sensors.base_camera import BaseCamera
        from metadrive.component.sensors.distance_detector import DistanceDetector

        parent, sensor_name = sensor_identifier.split(":")
        agent = self.agents[parent]
        sensor = self.sensors.get(sensor_name, None)

        if isinstance(sensor, BaseCamera):
            data = sensor.perceive(
                to_float=self.env_config["norm_pixel"],
                new_parent_node=agent.origin,
            )
            data = data[:, :, ::-1]
        elif isinstance(sensor, DistanceDetector):
            data, objs = sensor.perceive(
                agent,
                self.dynamic_world if "lidar" in sensor_name else self.static_world,
                agent.config[sensor_name]["num_lasers"],
                agent.config[sensor_name]["distance"],
                height=1.0 if "lidar" in sensor_name else None,
            )
            data = np.array(data)
        elif sensor_name in ["top_down", "topdown", "bev", "birdview"]:
            self.env.main_camera.track(agent)
            data = self.env.render(
                mode=sensor_name,
                screen_size=(500, 400),
                window=False,
            )
        else:
            raise ValueError(f"Unknown sensor type: {sensor_name}")

        return data

    @override(AdapterInterface)
    def get_actors(self):
        return list(self.agents.keys())

    @override(AdapterInterface)
    def get_sensors(self):
        actors = self.get_actors()
        sensors = list(self.sensors.keys()) + ["topdown"]
        return [f"{actor}:{sensor}" for actor in actors for sensor in sensors]

    @override(AdapterInterface)
    def load_scenario(self, scenario: dict = None):
        if scenario and scenario != self.config:
            self.cleanup()
            self.config = scenario
            self.initialize()

        self._last_obs, self._last_info = self.env.reset()
        self._agent_states = {
            agent_id: agent.get_state() for agent_id, agent in self.agents.items()
        }

    @override(AdapterInterface)
    def apply_control(self, target: str, control: dict):
        self._action_dict[target] = control

    @override(AdapterInterface)
    def tick(self):
        if len(self._action_dict) == len(self.agents):
            (
                self._last_obs,
                self._last_reward,
                self._last_terminated,
                self._last_truncated,
                self._last_info,
            ) = self.env.step(self._action_dict)

            self._action_dict.clear()
            self._agent_states = {
                agent_id: agent.get_state() for agent_id, agent in self.agents.items()
            }
            self._tick_cnt += 1

        return self._tick_cnt

    def get_state(self, agent: str):
        return self._agent_states[agent]
